package com.wyz.socket.core.impl.stealing;

import com.wyz.socket.core.base.IoTask;

import java.io.Closeable;
import java.io.IOException;
import java.nio.channels.*;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public abstract class IoStealingSelectorThread extends Thread implements Closeable {
    private static final int MAX_ONCE_READ_TASK = 128;
    private static final int MAX_ONCE_WRITE_TASK = 32;
    private static final int MAX_ONCE_RUN_TASK = MAX_ONCE_READ_TASK + MAX_ONCE_WRITE_TASK;

    // read 1 << 0  == 1 ;
    // write 1 << 2 == 100 = 4 ;
    // 001 |100 = 101 = 5 ;

    private static final int VALID_OPS = SelectionKey.OP_READ | SelectionKey.OP_WRITE;

    private final Selector selector;

    private volatile boolean isRunning = true;

    private final ArrayBlockingQueue<IoTask> readyTaskQueue = new ArrayBlockingQueue<>(MAX_ONCE_RUN_TASK);

    private final ConcurrentLinkedQueue<IoTask> registerTaskQueue = new ConcurrentLinkedQueue<>();

    private final AtomicBoolean unregisterLocker = new AtomicBoolean(false);

    private final List<IoTask> onceReadyReadTaskCache = new ArrayList<>(MAX_ONCE_READ_TASK);
    private final List<IoTask> onceReadyWriteTaskCache = new ArrayList<>(MAX_ONCE_WRITE_TASK);

    public IoStealingSelectorThread(String name, Selector selector) {
        super(name);
        this.selector = selector;
    }

    private final AtomicLong busyStatus = new AtomicLong();

    private StealingService stealingService;

    public void register(IoTask task) {
        if ((task.ops & ~VALID_OPS) != 0) {
            throw new UnsupportedOperationException("");
        }
        registerTaskQueue.offer(task);
        selector.wakeup();
    }

    public void unregister(SocketChannel channel) {
        SelectionKey selectionKey = channel.keyFor(selector);
        if (selectionKey != null && selectionKey.attachment() != null) {
            selectionKey.attach(null);

            if (Thread.currentThread() == this) {
                selectionKey.cancel();
            } else {
                synchronized (unregisterLocker) {
                    unregisterLocker.set(true);
                    selector.wakeup();
                    selectionKey.cancel();
                    unregisterLocker.set(false);
                }
            }
        }
    }

    private void offerReadyTaskQueue(final Queue<IoTask> readyQueue, final List<IoTask> readyList) {
        readyQueue.addAll(readyList);
    }

    private void consumeRegisterTasks(final Queue<IoTask> registerTaskQueue) {
        final Selector selector = this.selector;
        IoTask registerTask = registerTaskQueue.poll();
        while (registerTask != null) {
            try {
                final SocketChannel channel = registerTask.channel;
                int registerOps = registerTask.ops;
                // read == 1  write == 4
                //     001          100
                // VALID_OPS   101  > ~101 = 010 ;
                // read & VALID_OPS = 0 | write & VALID_OPS = 0

                SelectionKey selectionKey = channel.keyFor(selector);
                if (selectionKey == null) {
                    selectionKey = channel.register(selector, registerOps, new KeyAttachment());
                } else {
                    selectionKey.interestOps(selectionKey.interestOps() | registerOps);
                }
                Object attachment = selectionKey.attachment();
                if (attachment instanceof KeyAttachment) {
                    ((KeyAttachment) attachment).attach(registerOps, registerTask);
                } else {
                    selectionKey.cancel();
                }
            } catch (ClosedChannelException | CancelledKeyException | ClosedSelectorException ignored) {
                registerTask.onThrowException(ignored);
            } finally {
                registerTask = registerTaskQueue.poll();
            }
        }
    }

    private void consumeReadyTasks(final Queue<IoTask> readyTaskQueue, final Queue<IoTask> registerTaskQueue) {
        // 凯斯消费已就绪的读写事件
        final AtomicLong busyStatus = this.busyStatus;
        IoTask doTask = readyTaskQueue.poll();
        while (doTask != null) {
            busyStatus.incrementAndGet();
            if (processTask(doTask)) {
                registerTaskQueue.offer(doTask);
            }
            doTask = readyTaskQueue.poll();
        }
        // 消费完成进行窃取
        final StealingService stealingService = this.stealingService;
        if (stealingService != null) {
            doTask = stealingService.steal(readyTaskQueue);
            while (doTask != null) {
                busyStatus.incrementAndGet();
                if (processTask(doTask)) {
                    registerTaskQueue.offer(doTask);
                }
                doTask = stealingService.steal(readyTaskQueue);
            }
        }
    }


    @Override
    public void run() {
        super.run();
        final Selector selector = this.selector;
        final ArrayBlockingQueue<IoTask> readyTaskQueue = this.readyTaskQueue;
        final ConcurrentLinkedQueue<IoTask> registerTaskQueue = this.registerTaskQueue;
        final List<IoTask> onceReadyReadTaskCache = this.onceReadyReadTaskCache;
        final List<IoTask> onceReadyWriteTaskCache = this.onceReadyWriteTaskCache;
        final AtomicBoolean unregisterLocker = this.unregisterLocker;

        try {
            while (isRunning) {
                consumeRegisterTasks(registerTaskQueue);

                int select = selector.select();

                while (unregisterLocker.get()) {
                    Thread.yield();
                }

                if (select == 0) {
                    continue;
                }

                int onceReadTaskCacheSize = MAX_ONCE_READ_TASK;
                int onceWriteTaskCacheSize = MAX_ONCE_WRITE_TASK;
                Iterator<SelectionKey> iterator = selector.selectedKeys().iterator();
                while (iterator.hasNext()) {
                    SelectionKey selectionKey = iterator.next();
                    Object attachmentObj = selectionKey.attachment();
                    if (selectionKey.isValid() && attachmentObj instanceof KeyAttachment) {
                        KeyAttachment keyAttachment = (KeyAttachment) attachmentObj;
                        try {
                            int readyOps = selectionKey.readyOps(); //获取已经准备好的事件 读取或者写入
                            int interestOps = selectionKey.interestOps(); // 获取注册关注的事件 读取或者写入

                            if ((readyOps & SelectionKey.OP_READ) != 0 && onceReadTaskCacheSize-- > 0) { // 如果读取准备好
                                onceReadyReadTaskCache.add(keyAttachment.readableTask); // 添加上读取任务
                                interestOps = interestOps & ~SelectionKey.OP_READ; // 并取消读取关注 否则 容易重复 监听注册事件 造成重读生成 就绪事件
                            }

                            if ((readyOps & SelectionKey.OP_WRITE) != 0 && onceWriteTaskCacheSize-- > 0) {
                                onceReadyWriteTaskCache.add(keyAttachment.writeableTask);
                                interestOps = interestOps & ~SelectionKey.OP_WRITE;
                            }

                            selectionKey.interestOps(interestOps);
                        } catch (CancelledKeyException cancelledKeyException) {
                            if (keyAttachment.readableTask != null) {
                                onceReadyReadTaskCache.remove(keyAttachment.readableTask);
                            }
                            if (keyAttachment.writeableTask != null) {
                                onceReadyWriteTaskCache.remove(keyAttachment.writeableTask);
                            }
                        }
                    }
                    iterator.remove();
                }
                if (!onceReadyReadTaskCache.isEmpty()) {
                    offerReadyTaskQueue(readyTaskQueue, onceReadyReadTaskCache);
                    onceReadyReadTaskCache.clear();
                }

                if (!onceReadyWriteTaskCache.isEmpty()) {
                    offerReadyTaskQueue(readyTaskQueue, onceReadyWriteTaskCache);
                    onceReadyWriteTaskCache.clear();
                }

                consumeReadyTasks(readyTaskQueue, registerTaskQueue);
            }
        } catch (IOException e) {
            try {
                selector.close();
            } catch (IOException ioException) {
                ioException.printStackTrace();
            }
        } finally {
            readyTaskQueue.clear();
            registerTaskQueue.clear();
            onceReadyReadTaskCache.clear();
            onceReadyWriteTaskCache.clear();
        }
    }

    protected abstract boolean processTask(IoTask task);

    // 窃取


    public Queue<IoTask> getReadyTaskQueue() {
        return readyTaskQueue;
    }

    public Queue<IoTask> getRegisterTaskQueue() {
        return registerTaskQueue;
    }

    public void close() {
        isRunning = false;
        try {
            selector.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        interrupt();
    }

    public long getBusyStatus() {
        if (selector.isOpen()) {
            return busyStatus.get();
        } else {
            return -1L;
        }
    }

    public void setStealingService(StealingService stealingService) {
        this.stealingService = stealingService;
    }

    /**
     * 注册时添加的附件
     */
    public static class KeyAttachment {
        IoTask readableTask;
        IoTask writeableTask;

        void attach(int ops, IoTask task) {
            if (ops == SelectionKey.OP_READ) {
                readableTask = task;
            } else if (ops == SelectionKey.OP_WRITE) {
                writeableTask = task;
            }
        }
    }
}
